import os
import pandas as pd
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import torchvision.transforms as transforms

class MNISTDataset(Dataset):
    def __init__(self, img_dir, labels_file, transform=None):
        self.img_dir = img_dir
        self.labels_df = pd.read_excel(labels_file)
        self.transform = transform

    def __len__(self):
        return len(self.labels_df)

    def __getitem__(self, idx):
        img_name = self.labels_df.iloc[idx, 0] + '.jpg'
        label = self.labels_df.iloc[idx, 1]
        img_path = os.path.join(self.img_dir, img_name)
        image = Image.open(img_path).convert('L')  # Convert image to grayscale
        if self.transform:
            image = self.transform(image)
        return image, label

# Define your transformation with data augmentation: resize, random rotation, random crop, and normalization
transform = transforms.Compose([
    transforms.Resize((28, 28)),
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # MNIST normalization values
])


dataset = MNISTDataset(img_dir='./data/train_data', labels_file='./data/train_data_label.xlsx', transform=transform)
data_loader = DataLoader(dataset, batch_size=64, shuffle=True)

import torch
from CNN import SimpleCNN
import torch.nn as nn
model = SimpleCNN()

import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)


def train(model, device, train_loader, optimizer, criterion, epochs=10):
    model.train()
    for epoch in range(epochs):
        running_loss = 0.0
        correct = 0
        total = 0
        for batch_idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            loss.backward()
            optimizer.step()

            running_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()

        print(f'Epoch {epoch + 1}: Loss: {running_loss / len(train_loader)}, Accuracy: {100. * correct / total}%')


device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

train(model, device, data_loader, optimizer, criterion)

torch.save(model.state_dict(), 'mnist_cnn.pth')
print('Model saved to mnist_cnn.pth')
